import argparse
import os.path

import numpy as np


def compare(x, y, atol=1e-3, rtol=0, save_name_prefix=""):
    print(">> x.shape:", x.shape, "y.shape:", y.shape)
    z = x - y
    result = z[np.abs(z) > atol]
    print(">> diff result:\n", result)
    diff_save_name = save_name_prefix + "_diff_result.npy"
    np.save(diff_save_name, np.array(result))
    print("save diff result:", diff_save_name)

    indices = np.where(np.abs(z) > atol)
    print(">> diff num:", np.array(indices).shape)
    print(">> diff indices:\n", np.array(indices))
    indices_save_name = save_name_prefix + "_diff_indices.npy"
    np.save(indices_save_name, np.array(indices))
    print("save diff indices:", indices_save_name)

    print(">> allclose:")
    np.testing.assert_allclose(x, y, rtol=rtol, atol=atol, err_msg="ERROR===")


def load_npy(x_path, y_path):
    x = np.load(x_path)
    y = np.load(y_path)
    return x, y


def compare_diff(x_path, y_path):
    print("start to compare diff between:{0},{1}".format(x_path, y_path))
    x, y = load_npy(x_path, y_path)
    compare(x, y, save_name_prefix=os.path.splitext(x_path)[0])


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--ms_path', default=None, required=True, help='mindspore npy path')
    parser.add_argument('--pt_path', default=None, required=True, help="python npy path")
    args = parser.parse_args()
    ms_path = args.ms_path
    pt_path = args.pt_path
    print("ms_path:{0}".format(ms_path))
    print("pt_path:{0}".format(pt_path))
    compare_diff(ms_path, pt_path)
